import os

from langchain.llms.baidu_qianfan_endpoint import QianfanLLMEndpoint
from config import QIANFAN_AK, QIANFAN_SK

os.environ["QIANFAN_AK"] = QIANFAN_AK
os.environ["QIANFAN_SK"] = QIANFAN_SK

# 默认模型为 "ERNIE-Bot-turbo"
# temperature 用温度来控制输出结果的确定性：温度越高，越不确定；温度越低，越确定。
# 增大一点，设为较热的1，跟LLM说5次“你好，今天天气真不错”，它每次的回答都不相同：
# 把温度设为冷冷的0，5次回答有4次都一模一样，确定性非常好：
qianfan_model = QianfanLLMEndpoint(model="ERNIE-Bot",
                                   qianfan_ak="KBe3HbStsjvnGS9KFrNZFPes",
                                   qianfan_sk="3yGoXGIdCCy8QLwcLCLdh3UGGlM7PKNn",
                                   streaming=True,
                                   request_timeout=60,
                                   top_p=0.95,
                                   temperature=1,
                                   penalty_score=1)
# res = qianfan_model("你好，今天天气真不错")
# print(res)
#
# res = qianfan_model("你好，今天天气真不错")
# print(res)
#
# res = qianfan_model("你好，今天天气真不错")
# print(res)
#
# res = qianfan_model("你好，今天天气真不错")
# print(res)
#
# res = qianfan_model("你好，今天天气真不错")
# print(res)
#
# qianfan_model.predict('你好，今天天气真不错')


""" LLM 设计分析 """
"""BaseLanguageModel"""
# BaseLanguageModel 是父类，定义了方法
# generate_prompt / agenerate_prompt (同步\异步)
# predict / apredict
# predict_messages / apredict_messages
# get_token_ids
# get_num_tokens
# get_num_tokens_from_messages

"""BaseLLM"""
# BaseLLM 继承了 BaseLanguageModel
# invoke -> generate_prompt -> generate -> _generate  / ainvoke -> agenerate_prompt -> agenerate (同步\异步)
# batch -> generate_prompt -> generate -> _generate
# stream -> (invoke | _stream | ) ->
# _generate_helper -> _generate
# __call__ -> generate
# predict -> __call__ -> generate
# predict_messages -> __call__ -> generate

"""LLM"""
# LLM 继承了 BaseLLM
# 实现了 _generate 方法 -> _call
# 定义了抽象方法 _call 方法

"""QianfanLLMEndpoint"""
# QianfanLLMEndpoint 继承了 LLM
# 实现了 _call
# 实现了 _stream

# 上面的例子，无论是llm('Hi') (其调用的是 QianfanLLMEndpoint 父类的 _call_ 函数)
# 还是llm.predict('Hi, this a good day!')，其背后调用的都是 generate() 函数。
# 真实的调用关系是：predict() -> BaseLLM.__call__ -> BaseLLM.generate() -> _generate -> _call

"""generate() 函数"""
# generate() 函数允许传入一个 prompt 的list，返回一个 LLMResult
llm_result = qianfan_model.generate(['讲一个笑话', "作一首诗"])
# LLMResult
print(type(llm_result))
# [Generation(text='好的，让我来分享一个有趣的笑话：\n\n有一天，一只海龟在沙滩上散步，看到一只螃蟹在爬行。\n\n海龟问螃蟹：“你为什么总是横着走？”\n\n螃蟹回答说：“因为我有钳，所以任性。”')]
# [Generation(text='惜春\n春光易逝恨难留，花落无情水自流。\n莫待明年花更好，今朝有酒今朝醉。')]
for generation in llm_result.generations:
    print(generation)

# 返回了llm_output，包含输入和输出的token数量：
# 千帆目前返回是 None
print(llm_result.llm_output)

# 可以通过get_num_tokens()函数评估输入prompt的token数量（需要安装 tiktoken 包）
# 获取中文的Token，openai 一个英语单词1个token，一个中文字2个token
# 千帆目前获取不到
print(qianfan_model.get_num_tokens('讲一个笑话'))



